import torch
from d2l import torch as d2l

# 前面的例子中，我们仅展示了单个输入和单个输出通道的简化例子；这使得我们可以将输入、卷积核和输出看作二维张量。当我们添加通道时，我们的输入和隐
# 藏的表示都变成了三维张量。例如，每个RGB输入图像具有3 × h × w的 形状。我们将这个大小为3的轴称为通道（channel）维度。
# 为了加深理解，我们实现一下多输入通道互相关运算。简而言之，我们所做的就是对每个通道执行互相关操X作，然后将结果相加。
def corr2d_multi_in(X, K):
    return sum(d2l.corr2d(x, k) for x, k in zip(X, K))  # # 使用zip()函数将两个可迭代对象打包，并在循环中同时迭代

X = torch.tensor([[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]])
K = torch.tensor([[[0.0, 1.0], [2.0, 3.0]], [[1.0, 2.0], [3.0, 4.0]]])
corr2d_multi_in(X, K)

# 多输出通道，如下所示，我们实现一个计算多个通道的输出的互相关函数。
def corr2d_multi_in_out(X, K):
    # 迭代"K"的第0个维度，每次都对输入"X"执行互相关运算。
    # 最后将所有结果都叠加在一起
    return torch.stack([corr2d_multi_in(X, k) for k in K], 0)
# 创建一个3输出通道的卷积核
K = torch.stack((K, K + 1, K + 2), 0)
K.shape
corr2d_multi_in_out(X, K)

# 1x1卷积层
# 下面，我们使用全连接层实现1 × 1卷积。请注意，我们需要对输入和输出的数据形状进行调整。
def corr2d_multi_in_out_1x1(X, K):
    c_i, h, w = X.shape
    c_o = K.shape[0]
    X = X.reshape((c_i, h * w))
    K = K.reshape((c_o, c_i))
    # 全连接层中的矩阵乘法
    Y = torch.matmul(K, X)
    return Y.reshape((c_o, h, w))
# 当执行1 × 1卷积运算时，上述函数相当于先前实现的互相关函数corr2d_multi_in_out。让我们用一些样本 数据来验证这一点。
X = torch.normal(0, 1, (3, 3, 3))  # 使用torch.normal的话记得声明均值和标准差，即这里的0和1
K = torch.normal(0, 1, (2, 3, 1, 1))
Y1 = corr2d_multi_in_out_1x1(X, K)
Y2 = corr2d_multi_in_out(X, K)
assert float(torch.abs(Y1 - Y2).sum()) < 1e-6

# 小结
# • 多输入多输出通道可以用来扩展卷积层的模型。
# • 当以每像素为基础应用时，1 × 1卷积层相当于全连接层。
# • 1 × 1卷积层通常用于调整网络层的通道数量和控制模型复杂性。